{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Retail Object Recognition\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n",
    "\n",
    "## Metric Learning Recognition\n",
    "\n",
    "Retail Object Recognition pretrained models uses `metric learning recognition` pipeline in TAO. It is a classifier that encodes the input image to embedding vectors and predicts their labels based on the embedding vectors in the reference space. MLRecogNet consists of two parts\n",
    "\n",
    "* Trunk: A backbone network that encodes the input image to a feature vector.\n",
    "* Embedder: A fully connected layer that maps the feature vector to the embedding space.\n",
    "\n",
    "The embedding space is a high-dimensional space where the distance between the embedding vectors of the same class is small and the distance between the embedding vectors of different classes is large. The embedder is trained to minimize the distance between the embedding vectors of the same class and maximize the distance between the embedding vectors of different classes. The embedding vectors of the query images are compared with the embedding vectors of the reference images to predict the labels of the query images.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this Notebook, you learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Train a model for object recogtion on an [ImageNet](https://www.image-net.org/) format example retail classification dataset.\n",
    "* Evaluate the trained model & export results.\n",
    "* Run Inference on the trained model.\n",
    "* Export the trained model to an ONNX file for deployment to DeepStream or TensorRT.\n",
    "\n",
    "At the end of this Notebook, you have a generated, trained `MLRecog` model which you can deploy using [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This Notebook shows an example usecase of MLRecogNet using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Inferences](#head-6)\n",
    "7. [Deploy](#head-7)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the Docker, they must be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Use it to configure your specific directories so that data, specs, results, and cache directories are correctly visible to the Docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"ml_recognition\", \"results\")\n",
    "os.environ[\"HOST_MODEL_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"ml_recognition\", \"models\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=/path/to/local/tao-experiments/metric_learning_recognition\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR\n",
    "! mkdir -p $HOST_MODEL_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tao_configs = {\n",
    "   \"Mounts\":[\n",
    "       # Mapping the data directory\n",
    "       {\n",
    "           \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "           \"destination\": \"/workspace/tao-experiments\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_MODEL_DIR\"],\n",
    "           \"destination\": \"/model\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         }\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tao_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a Python package distributed as a Python wheel listed in PyPI. Install the launcher by executing the following cell.\n",
    "\n",
    "TAO Toolkit recommends running the TAO launcher in a virtual env with Python 3.6.9. Follow the [instructions](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a Python virtual env using the `virtualenv` and `virtualenvwrapper` packages. After you have setup the virtualenvwrapper, set the version of Python in the virtual env with the `VIRTUALENVWRAPPER_PYTHON` variable, by running:\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the Notebook from the virtual environment. In addition to installing TAO python package, you must meet the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "After you have installed the pre-requisites, log in to the Docker registry nvcr.io using the following command:\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "Enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.6.9, following these [instructions](https://virtualenvwrapper.readthedocs.io/en/latest/install.html), to set up a Python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 Prepare dataset <a class=\"anchor\" id=\"head-2\"></a>\n",
    "\n",
    "Here we use [Retail Product Checkout dataset](https://www.kaggle.com/datasets/diyer22/retail-product-checkout-dataset) to illustrate the method of training the metric recognition model for retail item recognition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Action required] Download the dataset manually.\n",
    "# [Action required] Put your downloaded .zip dataset file at $HOST_DATA_DIR/retail-product-checkout-dataset.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Extract the files\n",
    "# !apt-get install unzip\n",
    "!mkdir $HOST_DATA_DIR\n",
    "\n",
    "# set dataset root folder path\n",
    "%env DATA_FOLDER=retail-product-checkout-dataset_classification_demo\n",
    "\n",
    "# Run data processing script: \n",
    "# 1. crop the images and save as a classification dataset\n",
    "# 2. split the dataset as train/val/test/reference sets\n",
    "# 3. separate the classes to be known and unknown classes\n",
    "\n",
    "# install the pkgs needed for process script if needed\n",
    "!pip install Cython==0.29.36\n",
    "!pip install opencv-python\n",
    "!pip install pycocotools\n",
    "!pip install tqdm\n",
    "# now run the process script\n",
    "!python $NOTEBOOK_ROOT/process_retail_product_checkout_dataset.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Verify\n",
    "!ls -l $HOST_DATA_DIR/$DATA_FOLDER/known_classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l $HOST_DATA_DIR/$DATA_FOLDER/unknown_classes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 Prepare Pretrained Model\n",
    "\n",
    "We will use NGC CLI to get the pre-trained models. For more details, go to ngc.nvidia.com and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Action required] Uncomment below code if you didn't have NGC CLI yet\n",
    "# # Installing NGC CLI on the local machine.\n",
    "# ## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# # Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/retail_object_recognition:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Action required] Uncomment below code if you don't have pretrained model downloaded\n",
    "# # Pull pretrained model from NGC\n",
    "\n",
    "# # DINOv2-Large backbone and linear head (only available for NVAIE users)\n",
    "# # coming: download links\n",
    "\n",
    "# # FAN-Base backbone and linear head\n",
    "# !ngc registry model download-version nvidia/tao/pretrained_fan_classification_nvimagenet:fan_base_hybrid_nvimagenet --dest $HOST_MODEL_DIR/\n",
    "# !ngc registry model download-version nvidia/tao/retail_object_recognition:trainable_head_fan_base_v2.0 --dest $HOST_MODEL_DIR/\n",
    "\n",
    "# # ResNet101 backbone\n",
    "# !ngc registry model download-version nvidia/tao/trainable_v1.1:retail_object_recognition_v1.1.pth --dest $HOST_MODEL_DIR/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    " \n",
    "* results_dir: a global setup for output directories, would create train/evaluation/inference/export subdirectories based on subtasks. Can be overwritten by subtask ``results_dir`` fields.\n",
    "\n",
    "* model: configure the model setting\n",
    "  * backbone: type of backbone architecture, supported backbones: `resnet_50`, `resnet_101`, `fan_tiny`, `fan_small`, `fan_base`, `fan_large`, `nvdinov2_vit_large_legacy`\n",
    "  * pretrained_model_path: path for the pretrained model weights\n",
    "  * pretrained_trunk_path: path for trunk pretrained weights\n",
    "  * pretrained_embedder_path: path for embedder pretrained weights\n",
    "  * input_width: width of an input image\n",
    "  * input_height: height of an input image\n",
    "  * input_channels: number of color channels for input images, always in channel first format\n",
    "  * feat_dim: size of the output embedding\n",
    "\n",
    "* train: configure the training hyperparameters\n",
    "  * train_trunk: If false, the trunk parameters will be frozen. Default true.\n",
    "  * train_embedder: If false, the embedder parameters will be frozen. Default true.\n",
    "  * optim: configure optimizer\n",
    "  * num_epochs: number of epochs\n",
    "  * checkpoint_interval: enabling how often to store models\n",
    "  * grad_clip: enabling gradient clipping\n",
    "  * smooth_loss: enabling label smoothing feature, True/False\n",
    "  * batch_size: number of images in 1 batch for training\n",
    "  * val_batch_size: number of images in 1 batch for validation\n",
    "  * resume_training_checkpoint_path: resume .pth model training from a saved checkpoint\n",
    "  * report_accuracy_per_class: enabling accuracy per class report instead of average class accuracies, True/False\n",
    "  \n",
    "* dataset: configure the dataset and augmentation methods\n",
    "  * train_dataset: path for the train dataset directory\n",
    "  * val_dataset: map of the validation or test dataset directory. It contains reference and query set.\n",
    "  * workers: number of workers to do data loading\n",
    "  * pixel_mean: pixel mean in 3 channels for normalization\n",
    "  * pixel_std: pixel standard deviation in 3 channels for normalization\n",
    "  * prob: probability of randomly flipping images horizontally\n",
    "  * re_prob: constant for random erasing\n",
    "  * gaussian_blur: configurations for Gaussian blur\n",
    "  * color_augmentation: configurations for color augmentation\n",
    "  * num_instance: number of types 1 image is repeated in a batch\n",
    "  * class_map: path to the YAML file mapping dataset class name to the new class names\n",
    "\n",
    "* evaluate: configure evaluate subtask parameters\n",
    "  * checkpoint: the .pth model for evaluation\n",
    "  * trt_engine: path of the TensorRT engine for evaluate\n",
    "  * report_accuracy_per_class: enabling accuracy per class report instead of average class accuracies, True/False\n",
    "  * topk: get predictions by the k nearest neighbor\n",
    "  * batch_size: the batch size for evaluate\n",
    "  * results_dir: the evaluation output directory. Have priority over global `results_dir`\n",
    "\n",
    "* inference: configure inference subtask parameters\n",
    "  * inference_input_type: the format of query dataset, image/image_folder/classification_folder\n",
    "  * checkpoint: the .pth model for inference\n",
    "  * trt_engine: path of the TensorRT engine for inference\n",
    "  * input_path: the inference image/image folder/classification dataset folder\n",
    "  * topk: get predictions by the k nearest neighbors\n",
    "  * batch_size: the batch size for inference\n",
    "  * results_dir: the inference output directory. Have priority over global `results_dir`\n",
    " \n",
    "* export: configure export subtask parameters\n",
    "  * checkpoint: the .pth model for export (to ONNX file)\n",
    "  * onnx_file: the exported ONNX model path. Have priority over the default ONNX name created from `export.results_dir`. \n",
    "  * gpu_id: the index of a single GPU only for export. default 0.\n",
    " \n",
    " \n",
    "* gen_trt_engine: configure TensorRT generation subtask parameters\n",
    "  * gpu_id: the index of a single GPU only for TensorRT engine generation. default 0.\n",
    "  * onnx_file: path of the ONNX file that TensorRT engine converted from\n",
    "  * trt_engine: path of the TensorRT engine to generate\n",
    "  * batch_size: the batch size of the TensorRT engine. When `batch_size=-1`, a dynamic batch size trt engine would be generated.\n",
    "  * verbose: If True, verbose information of TensorRT generation would be printed out\n",
    "  * tensorrt: trt engine generation setup\n",
    "  * results_dir: the trt engine generation output folder. Have priority over global `results_dir`\n",
    "\n",
    "See the [TAO documentation - Metric Learning Recognition](https://docs.nvidia.com/tao/tao-toolkit/text/metric_learning_recognition/metric_learning_recognition.html) to get all the parameters that are configurable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/train.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above example is used to train nvdinov2_vit_large backbone (trunk). To switch to other pretrained models with different backbones, you may want to refer to configurations at `$HOST_SPECS_DIR/train_resnet.yaml`, `$HOST_SPECS_DIR/train_fan.yaml`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Action Required] Update `model.pretrained_model_path`/`model.pretrained_embedder_path`/`model.pretrained_trunk_path` if you want to try pretrained models"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models.\n",
    "* WARNING: Training will take several hours or one day to complete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env MODEL_DIR = /model\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Train Metric Learning Recognition model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will train a MLRecog model with ResNet101 backbone and 2048 embedding size output. The backbone would be loaded with weights trained by NVImageNetV2 (same classes as [ImageNet](https://www.image-net.org/) but using licensed datasets)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Train model\")\n",
    "! tao model ml_recog train \\\n",
    "              -e $SPECS_DIR/train.yaml \\\n",
    "              results_dir=$RESULTS_DIR \\\n",
    "              dataset.train_dataset=$DATA_DIR/$DATA_FOLDER/known_classes/train \\\n",
    "              dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "              dataset.val_dataset.query=$DATA_DIR/$DATA_FOLDER/known_classes/val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "## Training command for multi-gpu training. We can define the number of gpus and specify which GPU's are to be used by setting the `train.gpu_ids` parameter.\n",
    "## The following command will trigger multi-gpu training on gpu 0 and gpu 1.\n",
    "# ! tao model ml_recog train \\\n",
    "#               -e $SPECS_DIR/train.yaml \\\n",
    "#               results_dir=$RESULTS_DIR \\\n",
    "#               dataset.train_dataset=$DATA_DIR/$DATA_FOLDER/known_classes/train \\\n",
    "#               dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "#               dataset.val_dataset.query=$DATA_DIR/$DATA_FOLDER/known_classes/val \\\n",
    "#               train.gpu_ids=[0,1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"train/dino_model_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/train/retail_object_recognition_model.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/retail_object_recognition_model.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>\n",
    "Evaluate trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate on known classes\n",
    "! tao model ml_recog evaluate \\\n",
    "            -e $SPECS_DIR/evaluate.yaml \\\n",
    "            results_dir=$RESULTS_DIR \\\n",
    "            evaluate.checkpoint=$RESULTS_DIR/train/retail_object_recognition_model.pth \\\n",
    "            dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "            dataset.val_dataset.query=$DATA_DIR/$DATA_FOLDER/known_classes/test \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate on unknown classes\n",
    "! tao model ml_recog evaluate \\\n",
    "            -e $SPECS_DIR/evaluate.yaml \\\n",
    "            evaluate.results_dir=$RESULTS_DIR/eval_unknown \\\n",
    "            evaluate.checkpoint=$RESULTS_DIR/train/retail_object_recognition_model.pth \\\n",
    "            dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/unknown_classes/reference \\\n",
    "            dataset.val_dataset.query=$DATA_DIR/$DATA_FOLDER/unknown_classes/test \n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Inferences <a class=\"anchor\" id=\"head-6\"></a>\n",
    "In this section, we run the metric_learning inference tool to generate inferences with the trained models and save the results under `$RESULTS_DIR`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# run inference on known classes\n",
    "! tao model ml_recog inference \\\n",
    "                    -e $SPECS_DIR/infer.yaml \\\n",
    "                    results_dir=$RESULTS_DIR \\\n",
    "                    inference.checkpoint=$RESULTS_DIR/train/retail_object_recognition_model.pth \\\n",
    "                    dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "                    inference.input_path=$DATA_DIR/$DATA_FOLDER/known_classes/test "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run inference on unknown classes\n",
    "! tao model ml_recog inference \\\n",
    "                    -e $SPECS_DIR/infer.yaml \\\n",
    "                    inference.results_dir=$RESULTS_DIR/inference_unknown \\\n",
    "                    inference.checkpoint=$RESULTS_DIR/train/retail_object_recognition_model.pth \\\n",
    "                    dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/unknown_classes/reference \\\n",
    "                    inference.input_path=$DATA_DIR/$DATA_FOLDER/unknown_classes/test "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Inference results:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/inference\n",
    "!ls -lth $HOST_RESULTS_DIR/inference_unknown"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Deploy <a class=\"anchor\" id=\"head-7\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Export the model to ONNX model.\n",
    "! tao model ml_recog export \\\n",
    "                   -e $SPECS_DIR/export.yaml \\\n",
    "                   results_dir=$RESULTS_DIR \\\n",
    "                   export.checkpoint=$RESULTS_DIR/train/retail_object_recognition_model.pth \\\n",
    "                   export.onnx_file=$RESULTS_DIR/export/retail_object_recognition_model.onnx\n",
    "               "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate FP16 TensorRT engine using tao deploy\n",
    "# ResNet backbones support INT8 as well. But NVDINOv2 and FAN backbones does not.\n",
    "!tao deploy ml_recog gen_trt_engine -e $SPECS_DIR/gen_trt_engine.yaml \\\n",
    "                                       gen_trt_engine.onnx_file=$RESULTS_DIR/export/retail_object_recognition_model.onnx \\\n",
    "                                       gen_trt_engine.trt_engine=$RESULTS_DIR/gen_trt_engine/retail_object_recognition_model.engine \\\n",
    "                                       gen_trt_engine.tensorrt.data_type=fp16 \\\n",
    "                                       results_dir=$RESULTS_DIR\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Generated tensorrt engines and calibration files:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/gen_trt_engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate with generated TensorRT engine\n",
    "!tao deploy ml_recog evaluate -e $SPECS_DIR/evaluate.yaml \\\n",
    "                                 evaluate.trt_engine=$RESULTS_DIR/gen_trt_engine/retail_object_recognition_model.engine \\\n",
    "                                 results_dir=$RESULTS_DIR \\\n",
    "                                 dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "                                 dataset.val_dataset.query=$DATA_DIR/$DATA_FOLDER/known_classes/test\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference with generated TensorRT engine\n",
    "!tao deploy ml_recog inference -e $SPECS_DIR/infer.yaml \\\n",
    "                                  inference.trt_engine=$RESULTS_DIR/gen_trt_engine/retail_object_recognition_model.engine \\\n",
    "                                  results_dir=$RESULTS_DIR \\\n",
    "                                  dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "                                  inference.input_path=$DATA_DIR/$DATA_FOLDER/known_classes/test\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('TensorRT Inference results:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/trt_inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install pandas matplotlib\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import ast\n",
    "\n",
    "# Load the CSV file\n",
    "csv_file = os.path.join(os.environ['HOST_RESULTS_DIR'], \"trt_inference\", \"trt_result.csv\")\n",
    "df = pd.read_csv(csv_file, header=None, names=['image_path', 'prediction', 'distance'])\n",
    "\n",
    "# Convert string representations of lists to actual lists\n",
    "df['prediction'] = df['prediction'].apply(ast.literal_eval)\n",
    "df['distance'] = df['distance'].apply(ast.literal_eval)\n",
    "\n",
    "# Sample a few rows\n",
    "sampled_df = df.sample(n=5)  # Change the number 5 to sample more or fewer images\n",
    "\n",
    "# Visualize the images with predictions and distances\n",
    "for index, row in sampled_df.iterrows():\n",
    "    image_path = os.environ['HOST_DATA_DIR'] + row['image_path'][5:] # switch /data back to local data path\n",
    "    img = Image.open(image_path)\n",
    "    \n",
    "    plt.imshow(img)\n",
    "    plt.title(f\"Prediction: {row['prediction'][0]}, Distance: {row['distance'][0]:.4f}\")\n",
    "    plt.axis('off')\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
